#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-i', '--input', help='Input data file', default='ss_pgf.360x181x15.h0.nc')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
parser.add_argument('-l', '--loop', help='Loop the plot', action='store_true')
args = parser.parse_args()

fig = plt.figure(figsize=(16, 8))

zo = np.linspace(0, 12, 101) * units.km

u_levels = np.linspace(-0.12, 0.12, 13)
omg_levels = np.linspace(-0.0004, 0.0004, 9)
w_levels = np.linspace(-0.0004, 0.0004, 9)

def plot(data_path, fig):
	f = addfile(data_path).sel(lat=0, method='nearest')

	dt = parse_datetime(f.time[args.time_step]) - parse_datetime(f.time[0])
	fig.suptitle(f'Day {dt.days + 1} {int(dt.seconds / 3600):2d}:00', fontsize=16)

	zi = f.gz * units('m**2 / s**2') / g

	ax = fig.add_subplot(2, 2, 1)
	u = vinterp_z(zi[-1,:,:], f.u[-1,:,:], zo)
	plot_contour_lon(ax, u, cmap='bwr', levels=u_levels, left_string=f'U (m/s) min: {u.min().values:.2f}, max: {u.max().values:.2f}')

	ax = fig.add_subplot(2, 2, 2)
	omg = vinterp_z(zi[-1,:,:], f.omg[-1,:,:], zo)
	plot_contour_lon(ax, omg, cmap='bwr', levels=omg_levels, left_string=f'Omega (Pa/s) min: {omg.min().values:.2e}, max: {omg.max().values:.2e}')

	if 'w_lev' in f.variables:
		ax = fig.add_subplot(2, 2, 3)
		w = vinterp_z(f.gz_lev[-1,:,:] * units('m**2 / s**2') / g, f.w_lev[-1,:,:], zo)
		plot_contour_lon(ax, w, cmap='bwr', levels=w_levels, left_string=f'W (m/s) min: {w.min().values:.2e}, max: {w.max().values:.2e}')

	ax = fig.add_subplot(2, 2, 4)
	for k in range(f.lev.size):
		ax.plot(f.lon, f.gz.isel(lev=k)[-1,:] / g / 1000, linewidth=0.5, color='k')

plot(args.input, fig)
if args.output:
	plt.savefig(args.output, dpi=300)
elif args.loop:
	loop_plot(args.input, fig, plot)
else:
	plt.tight_layout()
	plt.show()
